import sys
import puzautils.table.table as pt
import puzautils.table.puza_row as pr
import puzautils.table.puza_stream_table as pst
import puzautils.table.puza_memory_table as pmt

target_table = '/Projects/WUME/lih307/project/personal_retweet/filter/target_tweets_id_text.db'
neighbor_table = '/Projects/WUME/lih307/project/personal_retweet/filter/neighbor_tweets_id_text.db'
target_retweets_naive = '/Projects/WUME/lih307/project/personal_retweet/filter/target_tweets_retweet_naive.db'
target_retweets_naive_none = '/Projects/WUME/lih307/project/personal_retweet/filter/target_tweets_retweet_naive_none.db'
target_retweets_naive_verify = '/Projects/WUME/lih307/project/personal_retweet/filter/target_tweets_retweet_naive_verify.db'
graph = '/Projects/WUME/lih307/db/wume_tweets/raw/graph.db'

def construct_graph():
    global graph
    Graph = pmt.PuzaMemoryTable()
    Graph.Read(graph)
    Graph.Load()
    g = {}
    graph_len = len(Graph)
    for i in range(graph_len):
        source_id = Graph[i]['source_id']
        if not source_id in g:
            g[source_id] = set()
        target_ids = Graph[i]['target_ids']
        for id in target_ids:
            g[source_id].add(id)
    return g

def verify_naive_retweets():
    global target_retweets_naive
    global target_retweets_naive_verify
    global neighbor_table
    Target = pmt.PuzaMemoryTable()
    Target.Read(target_retweets_naive)
    Target.Load()
    all_ids = {}
    target_len = len(Target)
    for i in range(target_len):
        all_ids[str(Target[i]['previous_id'])] = i
    ## start
    NeighborTable = pst.PuzaStreamTable()
    NeighborTable.Read(neighbor_table)
    OutputTable = pt.PuzaTable()
    OutputTable.Save(target_retweets_naive_verify)
    while True:
        raw_obj = NeighborTable.StreamInRow()
        if raw_obj != None:
            r = pr.PuzaRow.Load(raw_obj)
            id = str(r['id'])
            if id in all_ids:
                OutputTable.StreamOutRow(Target[all_ids[id]], pointer = 0, lformat = 0, block_number = 1)
    OutputTable.StreamClose()       

def naive_retweet(graph):
    global target_table
    global target_retweets_naive
    global target_retweets_naive_none
    TargetTable = pst.PuzaStreamTable()
    TargetTable.Read(target_table)
    OutputTable = pt.PuzaTable()
    OutputTable.Save(target_retweets_naive)
    OutputTable2 = pt.PuzaTable()
    OutputTable2.Save(target_retweets_naive_none)
    Total_Number = 0
    RT = 0
    while True:
        raw_obj =  TargetTable.StreamInRow()
        if raw_obj != None:
            r = pr.PuzaRow.Load(raw_obj)
            if r['previous_id'] != None:
                current_user_id = str(r['user_id'])
                previous_user_id = str(r['previous_user_id'])
                if current_user_id in graph:
                    if previous_user_id in graph[current_user_id]:
                        OutputTable.StreamOutRow(r, pointer = 0, lformat = 0, block_number = 1)
                        RT = RT + 1
                    else:
                        OutputTable2.StreamOutRow(r, pointer = 0, lformat = 0, block_number = 1)
                Total_Number += 1
        else:
            break
    print 'Total RT',Total_Number
    print 'Graph RT', RT
    OutputTable.StreamClose()
    OutputTable2.StreamClose()

if __name__ == '__main__':
    # ========== construct naive retweets step 1 ==========
    g = construct_graph()
    naive_retweet(g)
    # ========== construct naive retweets step 2 ==========
    # verify_naive_retweets()
